package com.yang.limit.service.impl;

import com.yang.limit.constant.Constant;
import com.yang.limit.constant.RequestLimitType;
import com.yang.limit.dto.RequestLimitDTO;
import com.yang.limit.service.RequestLimitService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
import org.springframework.stereotype.Service;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

@Slf4j
@Service
public class TokenBucketServiceImpl implements RequestLimitService {

    @Resource(name = "tokenBucketThreadPoolScheduler")
    private ThreadPoolTaskScheduler threadPoolTaskScheduler;

    @Value("${request-limit.scan-package}")
    private String scanPackage;

    @Autowired
    private ResourcePatternResolver resourcePatternResolver;

    @Autowired
    private RedisTemplate<String, Object> redisTemplate;


    private static final DefaultRedisScript<Long> TOKEN_BUCKET_SCRIPT;

    static {
        TOKEN_BUCKET_SCRIPT = new DefaultRedisScript<>();
        TOKEN_BUCKET_SCRIPT.setLocation(new ClassPathResource("token_bucket.lua"));
        TOKEN_BUCKET_SCRIPT.setResultType(Long.class);
    }


    @Override
    public RequestLimitType getLimitType() {
        return RequestLimitType.TOKEN_BUCKET;
    }

    @Override
    public boolean checkIsRequestLimit(RequestLimitDTO dto) {
        String key = Constant.REDIS_KEY_TOKEN_BUCKET_PREFIX + dto.getKey();
        List<String> params = new ArrayList<>();
        params.add(key);
        // 1.执行lua脚本
        Long result = redisTemplate.execute(
                TOKEN_BUCKET_SCRIPT,
                params);
        if (result == -1){
            log.info("【{}】限流控制策略，令牌桶中不存在令牌，当前请求被拦截", dto.getKey());
            return true;
        }
        else{
            log.info("【{}】令牌桶存在令牌，未达到限流值，放行当前请求", dto.getKey());
            return false;
        }
    }

    @PostConstruct
    public void addToken() {
        List<RequestLimitDTO> tokenLimitList = this.getTokenLimitList(resourcePatternResolver, RequestLimitType.TOKEN_BUCKET, scanPackage);
        //项目中没有使用到接口限流
        if (tokenLimitList.isEmpty()) {
            log.info("项目中没有使用接口限流");
            return;
        }
        //遍历接口
        tokenLimitList.forEach(task -> threadPoolTaskScheduler.scheduleAtFixedRate(() -> {
            String key = Constant.REDIS_KEY_TOKEN_BUCKET_PREFIX + task.getKey();
            Integer count = (Integer) redisTemplate.opsForValue().get(key);
            if (count == null) count = 0;
            if (count >= task.getLimiter().limitCount()) {
                log.info("【{}】令牌数量已达最大值【{}】，丢弃新生成令牌", key, count);
                return;
            }
            //添加令牌
            int remainSize = Math.min(task.getLimiter().limitPeriodCount(), task.getLimiter().limitCount());
            for (int i = 0; i < remainSize; i++) {
                redisTemplate.opsForValue().increment(key);
            }
            log.info("【{}】生成令牌丢入令牌桶，当前令牌数：{}，令牌桶容量：{}", key, count + remainSize, task.getLimiter().limitCount());
        }, task.getLimiter().period()));
    }
}
